#include "trie.h"

Trie::Trie()
{
    m_root = new TrieNode();
    m_size = 0;
}

Trie::~Trie()
{
    if (m_root != nullptr)
    {
        delete m_root;
        m_root = nullptr;
    }
    m_size = 0;
}

void Trie::insert(const std::string & key)
{
    insert(key.c_str(), key.size());
}

void Trie::insert(const char * data, int size)
{
    TrieNode * node = m_root;
    for (int i = 0; i < size; i++)
    {
        int idx;
        bool found = false;
        for (idx = 0; idx < node->m_size; idx++)
        {
            if (data[i] == node->m_children[idx]->m_data)
            {
                found = true;
                break;
            }
            else if (data[i] < node->m_children[idx]->m_data)
            {
                break;
            }
        }

        if (found)
        {
            node = node->m_children[idx];
        }
        else 
        {
            TrieNode * child = new TrieNode();
            child->m_data = data[i];
            node->insert(child, idx);
            node = child;
        }
        if (i == size - 1)
        {
            node->m_flag = M_FLAG_HAVE_STR;    // 不需要创建新的节点，插入的单词在trie里已经存在
        }
    }
    m_size += 1;
}

TrieNode * Trie::search(const char * data, int size)
{
    TrieNode * node = m_root;
    for (int i = 0; i < size; i++)
    {
        int idx;
        bool found = false;
        for (idx = 0; idx < node->m_size; idx++)
        {
            if (data[i] == node->m_children[idx]->m_data)
            {
                found = true;
                break;
            }
            else if (data[i] < node->m_children[idx]->m_data)
            {
                break;
            }
        }
        if (!found)
            return nullptr;
        node = node->m_children[idx];
    }
    return node;
}

bool Trie::prefix(const std::string & key)
{
    TrieNode * node = search(key.data(), key.size());
    return node != nullptr;
}

bool Trie::search(const std::string & key)
{
    TrieNode * node = search(key.data(), key.size());
    if (!node)
    {
        return false;
    }
    return node->m_flag;
}

void Trie::remove(const char * data, int size)
{
    auto node = m_root;
    TrieNode * last_branch_node = m_root;
    int last_branch_idx = 0;
    for (int i = 0; i < size; i++)
    {
        int idx;
        bool found = false;
        for (idx = 0; idx < node->m_size; idx++)
        {
            if (data[i] == node->m_children[idx]->m_data)
            {
                found = true;
                break;
            }
            else if (data[i] < node->m_children[idx]->m_data)
            {
                break;
            }
        }
        if (!found)
            return;
        if (node->m_size > 1 || node->m_flag == M_FLAG_HAVE_STR)
        {
            last_branch_node = node;
            last_branch_idx = idx;
        }
        node = node->m_children[idx];
    }

    if (node->m_size > 0)
    {
        // case 1
        node->m_flag = M_FLAG_NO_STR;
    }
    else
    {
        last_branch_node->remove(last_branch_idx);
    }
}

void Trie::remove(const std::string & key)
{
    remove(key.data(), key.size());
}

std::vector<std::string> Trie::getAllStr()
{
    std::queue<std::pair<std::string, TrieNode*>> q;
    std::vector<std::string> outStrs;

    // 入队操作
    std::pair<std::string, TrieNode*> p = std::make_pair("", m_root);
    q.push(p);

    while (!q.empty())
    {
        auto p2 = q.front();
        TrieNode * trieNode = p.second;
        for (int i = 0; i < trieNode->m_size; i++)
        {
            std::string nowStr = p2.first + trieNode->m_children[i]->m_data;
            p = std::make_pair(nowStr, trieNode->m_children[i]);
            q.push(p);
            if (trieNode->m_children[i]->m_flag == M_FLAG_HAVE_STR)
            {
                outStrs.push_back(nowStr);
            }
        }
        q.pop();
    }

    return outStrs;
}